import argparse
from tqdm import tqdm
import pickle
import random
from loguru import logger
from gensim.models import Word2Vec


def random_walk(graph, max_length=100):
    """Given a block graph (In that graph, all the node is BlockNode 
    and we use the addr to represent the node)

    Args:
        graph (Graph): A NetworkX Graph class
        max_length (int): Maximum length of a trace

    Returns:
        list: A trace of block addresses
    """
    first_block = list(graph.nodes)[0]
    trace = [first_block]
    while len(list(graph.successors(trace[-1]))) > 0 \
            and len(trace) < max_length:
        block_list = list(graph.successors(trace[-1]))
        # Remove blocks with 0 size
        # That can handle some special blocks
        # (e.g. Function UnresolvableCellTarget)
        new_list = []
        for b in block_list:
            if b.size > 0:
                new_list.append(b)
        block_list = new_list

        next_block = block_list[random.randint(0, len(block_list) - 1)]
        trace.append(next_block)

    block_trace = []
    for t in trace:
        block_trace.append(str(t.addr))
    return block_trace


def run(args):

    graph_data = None
    with open(args.graph_data, "rb") as f:
        graph_data = pickle.load(f)

    if graph_data is None:
        logger.error("No Graph Data !")
        exit(0)

    pbar = tqdm(total=len(graph_data.keys()))

    deep_walk_data = {}

    for filename in graph_data:
        pbar.update(1)
        deep_walk_data[filename] = {}
        for function_name in graph_data[filename]:

            deep_walk_data[filename][function_name] = {}
            g = graph_data[filename][function_name]
            trace = []
            for _ in range(args.random_times):
                trace.append(random_walk(g))

            model = Word2Vec(sentences=trace, vector_size=200, window=10, min_count=1, workers=4)
            for _, word in enumerate(model.wv.index_to_key):
                deep_walk_data[filename][function_name][word] = model.wv[word]

    pbar.close()

    logger.info("Saving the data ...")

    with open(args.output, "wb") as f:
        pickle.dump(deep_walk_data, f)
    logger.info("Finished!")


if __name__ == "__main__":
    parser = argparse.ArgumentParser("Generate the deep walk embedding from CFG")
    parser.add_argument("-o", "--output", type=str, default="deep_walk.pkl", help="The name of the model")
    parser.add_argument("-g", "--graph_data", type=str, default="None", help="The history of graph data")
    parser.add_argument("-rt", "--random_times", type=int, default=100, help="The times of the random walk")

    args = parser.parse_args()
    run(args)
